import time
import paramiko
import re
from robot.api.deco import keyword

@keyword("Try Connect SSH")
def try_connect_SSH(hostname, port, username, password):
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    
    try:
        # 尝试建立SSH连接
        ssh.connect(hostname, int(port), username, password, timeout=10)
        print("✅ ssh连接成功!")
        return ssh
    except Exception as e:
        print(f"❌ ssh连接失败: {str(e)}")
        return None

@keyword("Send Ssh Msg")
def send_ssh_msg(ssh, command):
    if not ssh:
        return "❌ SSH连接未建议!"
    
    try:
        shell = ssh.invoke_shell()
        shell.settimeout(5)
        shell.send(command + "\n")

        output = ""
        start_time = time.time()
        timeout = 30

        # 等待命令提示符出现
        prompt_pattern = re.compile(r'[$#]\s*$')

        while time.time() - start_time < timeout:
            time.sleep(0.5)
            if shell.recv_ready():
                chunk = shell.recv(4096).decode('utf-8', 'ignore')
                output += chunk

                # 检测命令提示符表示命令执行完成
                if prompt_pattern.search(chunk):
                    break
            else:
                time.sleep(0.1)
        
        return output.strip()
    except Exception as e:
        return f"❌ 命令执行错误: {str(e)}"

@keyword("Close SSH Connection")
def close_ssh_connection(ssh):
    if ssh:
        ssh.close()
        print("✅ SSH连接已关闭!")